"""
生成示例数据脚本
用于测试项目功能（生成模拟的雷达图像和标签）
"""
import os
import numpy as np
import cv2
from tqdm import tqdm


def generate_radar_image(size=256):
    """
    生成模拟的雷达俯视图
    包含旁瓣、条纹、弯曲带状噪声等
    """
    # 创建基础图像
    image = np.zeros((size, size), dtype=np.uint8)

    # 添加背景噪声
    noise = np.random.normal(20, 10, (size, size))
    image = np.clip(image + noise, 0, 255).astype(np.uint8)

    # 添加条纹噪声（模拟旁瓣）
    for i in range(0, size, 10):
        stripe_intensity = np.random.randint(30, 60)
        image[i:i+2, :] = np.clip(image[i:i+2, :] + stripe_intensity, 0, 255)

    # 添加弯曲带状噪声
    for _ in range(3):
        center_y = np.random.randint(size // 4, 3 * size // 4)
        amplitude = np.random.randint(10, 30)
        frequency = np.random.uniform(0.02, 0.05)

        for x in range(size):
            y_offset = int(amplitude * np.sin(frequency * x))
            y = center_y + y_offset

            if 0 <= y < size:
                width = np.random.randint(3, 8)
                intensity = np.random.randint(40, 80)
                y_start = max(0, y - width // 2)
                y_end = min(size, y + width // 2)
                image[y_start:y_end, x] = np.clip(image[y_start:y_end, x] + intensity, 0, 255)

    # 添加随机斑点噪声
    num_spots = np.random.randint(50, 100)
    for _ in range(num_spots):
        x = np.random.randint(0, size)
        y = np.random.randint(0, size)
        radius = np.random.randint(1, 3)
        intensity = np.random.randint(30, 100)
        cv2.circle(image, (x, y), radius, int(intensity), -1)

    return image


def generate_mask(size=256, num_targets=None):
    """
    生成干净的目标点标签
    """
    mask = np.zeros((size, size), dtype=np.uint8)

    # 随机生成目标点数量
    if num_targets is None:
        num_targets = np.random.randint(3, 10)

    for _ in range(num_targets):
        x = np.random.randint(20, size - 20)
        y = np.random.randint(20, size - 20)

        # 目标点大小
        radius = np.random.randint(2, 5)

        # 绘制目标点（白色）
        cv2.circle(mask, (x, y), radius, 255, -1)

    return mask


def generate_dataset(num_samples=100, output_dir='data'):
    """
    生成完整的数据集

    Args:
        num_samples: 生成的样本数量
        output_dir: 输出目录
    """
    images_dir = os.path.join(output_dir, 'images')
    masks_dir = os.path.join(output_dir, 'masks')

    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(masks_dir, exist_ok=True)

    print(f"生成 {num_samples} 个样本...")

    for i in tqdm(range(num_samples)):
        # 生成文件名（4位数字，补零）
        filename = f"{i:04d}.png"

        # 生成图像和标签
        image = generate_radar_image()
        mask = generate_mask()

        # 保存
        image_path = os.path.join(images_dir, filename)
        mask_path = os.path.join(masks_dir, filename)

        cv2.imwrite(image_path, image)
        cv2.imwrite(mask_path, mask)

    print(f"\n完成！数据已保存到:")
    print(f"  图像: {images_dir}")
    print(f"  标签: {masks_dir}")


def main():
    """主函数"""
    print("=" * 80)
    print("生成示例数据")
    print("=" * 80)
    print("\n注意：这是用于测试的模拟数据，不是真实的雷达数据")
    print("如果您有真实数据，请直接将其放入 data/images/ 和 data/masks/\n")

    # 检查是否已有数据
    if os.path.exists('data/images') and len(os.listdir('data/images')) > 0:
        print("检测到 data/images/ 中已有数据")
        overwrite = input("是否覆盖？(y/n): ").strip().lower()
        if overwrite != 'y':
            print("已取消")
            return

    # 询问生成数量
    print("\n建议生成数量：")
    print("  - 快速测试: 100 个样本")
    print("  - 完整测试: 1500 个样本（与实际数据集相同）")

    num_samples_input = input("\n请输入要生成的样本数量（默认100）: ").strip()

    if num_samples_input:
        try:
            num_samples = int(num_samples_input)
        except ValueError:
            print("输入无效，使用默认值100")
            num_samples = 100
    else:
        num_samples = 100

    # 生成数据
    generate_dataset(num_samples=num_samples)

    print("\n" + "=" * 80)
    print("数据生成完成！")
    print("\n接下来可以：")
    print("  1. 运行环境检查: python check_environment.py")
    print("  2. 开始训练: python quick_start.py")
    print("=" * 80 + "\n")


if __name__ == '__main__':
    main()
